"""
https://leetcode.cn/problems/word-squares
"""
from typing import List

class Node:
    def __init__(self):
        self.next = {}
        # self.end = False
        self.passed = []

class Solution:

    def add_word(self, word):
        cur = self.root
        for c in word:
            if c not in cur.next:
                cur.next[c] = Node()
            cur = cur.next[c]
            cur.passed.append(word)
        # cur.end = True

    def prefix(self, prefix:str):
        cur = self.root
        for c in prefix:
            if c not in cur.next:
                return []
            cur = cur.next[c]
        # res = []
        # path = prefix.split()
        # def dfs(cur: Node):
        #     if cur.end:
        #         res.append("".join(path))
        #     for c, child in cur.next.items():
        #         path.append(c)
        #         dfs(child)
        #         path.pop()
        # dfs(cur)
        # return res
        return cur.passed

    def wordSquares(self, words: List[str]) -> List[List[str]]:
        m, n = len(words), len(words[0])
        self.root = Node()
        for word in words:
            self.add_word(word)

        res = []
        def dfs(step):
            if step == n:
                res.append(word_square[:])
                return
            # print(step, word_square)
            prefix = ''.join(word[step] for word in word_square)
            can_chose_words = self.prefix(prefix)
            for word in can_chose_words:
                word_square.append(word)
                dfs(step + 1)
                word_square.pop()

        word_square = []
        for word in words:
            word_square.append(word)
            dfs(1)
            word_square.pop()
        return res

words = ["area", "lead", "wall", "lady", "ball"]
# words = ["area", "arac", "aaef", "lady", "ball"]
print(Solution().wordSquares(words))
